import torch
from torch.utils.data import DataLoader, TensorDataset, random_split, Dataset
from prework.preprocess import train_data
# from prework.preprocess_plus import train_data


def data_deal(batch_size: int, shuffle: bool = True, train_ratio=0.8, val_ratio=0.1, random: bool = True):
    dataset = GetDataset()
    dataset_len = len(dataset)
    train_size = int(train_ratio * dataset_len)
    val_size = int(val_ratio * dataset_len)
    test_size = dataset_len - train_size - val_size
    if random is True:
        train_dataset, val_dataset, test_dataset = random_split(dataset, [train_size, val_size, test_size])
    else:
        return DataLoader(dataset, batch_size=batch_size, shuffle=shuffle)
    if val_size == 0:
        return DataLoader(train_dataset, batch_size=batch_size, shuffle=shuffle), DataLoader(test_dataset, batch_size=batch_size)
    elif test_size == 0:
        return DataLoader(train_dataset, batch_size=batch_size, shuffle=shuffle), DataLoader(val_dataset, batch_size=batch_size)
    else:
        return DataLoader(train_dataset, batch_size=batch_size, shuffle=shuffle), DataLoader(val_dataset, batch_size=batch_size), DataLoader(test_dataset, batch_size=batch_size)


class GetDataset(Dataset):

    def __init__(self):
        self.input, self.label = train_data()

    def __getitem__(self, index):
        x = self.input[index]
        for i in range(len(x)):
            x[i] = float(x[i])
        y = [int(self.label[index])]
        return torch.FloatTensor(x), torch.LongTensor(y)

    def __len__(self):
        return len(self.label)
